from classify.bayes.basic import *

class MultinomialNB(NaiveBayes):

    def feed_data(self, x, y, sample_weight=None):

        if isinstance(x, list):
            features = map(list, zip(*x))
        else:
            features = x.T

        features = [set(feat) for feat in features]
        feat_dics = [{_l: i for i, _l in enumerate(feats)} for feats in features]
        label_dic = {_l: i for i, _l in enumerate(set(y))}

        x = np.array([[feat_dics[i][_l] for i, _l in enumerate(sample)] for sample in x])

        y = np.array([label_dic[yy] for yy in y])

        cat_counter = np.bincount(y)

        n_possiilities = [len(feats) for feats in features]

        labels = [y == value for value in range(len(cat_counter))]

        labelled_x = [x[ci].T for ci in labels]

        self._x, self._y = x, y
        self._labelled_x, self._label_zip = labelled_x, list(zip(labels, labelled_x))
        (self._cat_counter, self._feat_dics, self._n_possibilities) = (
            cat_counter, feat_dics, n_possiilities)
        self.label_dic = {i: _l for _l, i in label_dic.items()}

        self.feed_sample_weight(sample_weight)

    def feed_sample_weight(self, sample_weight=None):
        self._con_counter = []

        for dim, _p in enumerate(self._n_possibilities):
            if sample_weight is None:
                self._con_counter.append([
                    np.bincount(xx[dim], minlength=_p) for xx in self._labelled_x])
            else:
                self._con_counter.append([
                    np.bincount(xx[dim], weights=sample_weight[
                    label] / sample_weight[label].mean(), minlength=_p)
                    for label, xx in self._label_zip])

    def _fit(self,lb):
        n_dim = len(self._n_possibilities)
        n_category = len(self._cat_counter)
        p_category = self.get_prior_probability(lb)

        data = [None] * n_dim
        for dim, n_possibilities in enumerate(self._n_possibilities):
            data[dim] = [[
                (self._con_counter[dim][c][p] + lb) / (self._cat_counter[c] + lb * n_possibilities)
                for p in range(n_possibilities)
            ] for c in range(n_category)]
        self._data = [np.array(dim_info) for dim_info in data]

        def func(input_x, tar_category):
            rs = 1

            for d, xx in enumerate(input_x):
                rs *= data[d][tar_category][xx]

            return rs * p_category[tar_category]

        return func

    def _transfer_x(self, x):

        for j, char in enumerate(x):
            x[j] = self._feat_dics[j][char]
        return x